package trainer;

import java.util.*;
import java.io.*;

/**
 * 
 * @author Chung-Hau Liang
 *
 */
public class HeldOutTest {

  ArrayList<String> total_data = new ArrayList<String>();

  ArrayList<String> heldout_data = new ArrayList<String>();

  ArrayList<String> train_data = new ArrayList<String>();

  int nPartitions = 2;

  /**
   * Train the model iteratively and choose the best lambda
   * 
   * @param filename
   * @return
   */
  public double HeldOutTest(String filename) {
    double dPP = 0.0, dMinPP = 999.0, dLambda = 0.1, dMinLambda = 999.0, dSumPP = 0.0;

    loadTrainData(filename);
    // for each candidate value of lambda;
    for (int i = 0; i < 10; i++) {

      for (int j = 0; j < nPartitions; j++) {

        System.err.println("Processing lambda: " + String.valueOf(dLambda) + " partition: "
                + String.valueOf(j + 1));
        partitionFile(j + 1, nPartitions);
        TrainModelByFile tm = new TrainModelByFile("train_tmp.txt", "model_tmp.txt", dLambda);
        dPP += TestData("heldout_tmp.txt", "model_tmp.txt");

        // clean up
        train_data.clear();
        heldout_data.clear();
      }
      System.err.println("avg PP:" + String.valueOf(dPP / (double) nPartitions));
      dPP /= (double) nPartitions;
      if (dPP < dMinPP) {
        dMinPP = dPP;
        dMinLambda = dLambda;
      }
      dPP = 0.0;
      dLambda += 0.1;

    }
    System.out.println("Optimal lambda:" + String.valueOf(dMinLambda));
    return dMinLambda;
  }

  /**
   * Read training data from file
   * 
   * @param filename_data
   */
  private void loadTrainData(String filename_data) {
    // load data

    try {
      BufferedReader in = new BufferedReader(new FileReader(filename_data));
      String line = "";
      while ((line = in.readLine()) != null) {
        total_data.add(line);
      }
      in.close();
    } catch (IOException e) {
    }
  }

  /**
   * Partition the training data into tmp files
   * 
   * @param nP
   * @param nT
   */
  private void partitionFile(int nP, int nT) {
    // @Para filename: training data
    // @Para nP: Current id of partition
    // @Para nT: Total number of partitions

    int nPartition_Size = total_data.size() / nT;
    int nStart = (nP - 1) * nPartition_Size;
    int nEnd = nStart + nPartition_Size;
    if (nEnd > total_data.size())
      nEnd = total_data.size();

    for (int i = nStart; i < nEnd; i++) {
      heldout_data.add(total_data.get(i));
    }
    for (int i = 0; i < total_data.size(); i++) {
      if ((i < nStart) || (i > nEnd)) {
        train_data.add(total_data.get(i));
      }
    }

    try {
      PrintWriter out = new PrintWriter(new FileWriter("train_tmp.txt"));
      for (String s : train_data)
        out.println(s);
      out.close();
      out = new PrintWriter(new FileWriter("heldout_tmp.txt"));
      for (String s : heldout_data)
        out.println(s);
      out.close();
    } catch (IOException e) {
    }

  }

  /**
   * Test the model using the data from the file
   * 
   * @param filename_data
   * @param filename_model
   * @return
   */
  private double TestData(String filename_data, String filename_model) {
    double dPP = 0.0;
    HashMap<String, Double> model = new HashMap<String, Double>();
    // load model
    try {
      BufferedReader in = new BufferedReader(new FileReader(filename_model));
      String line = "";
      while ((line = in.readLine()) != null) {
        if ((line.indexOf("Unigram") < 0) && (line.indexOf("Bigram") < 0)
                && (line.indexOf("Trigram") < 0)) {
          String[] strModel = line.split("\t");
          model.put(strModel[0], Double.valueOf(strModel[1]));
        }
      }
      in.close();
    } catch (IOException e) {
    }
    // load data
    ArrayList<String> data = new ArrayList<String>();
    try {
      BufferedReader in = new BufferedReader(new FileReader(filename_data));
      String line = "";
      data.add("<S>");
      data.add("<S>");
      while ((line = in.readLine()) != null) {

        data.add(line);
        if (line.equals("<PERIOD>")) {
          data.add("<S>");
          data.add("<S>");
        }
      }
      in.close();
    } catch (IOException e) {
    }
    // compute Perplexity
    double dH = 0.0;
    String target = "";
    int numActualTags = 0;
    for (int i = 2; i < data.size(); i++) {
      target = data.get(i - 2) + " " + data.get(i - 1) + " " + data.get(i);
      if (data.get(i).equalsIgnoreCase("<S>"))
        continue;
      double dP = model.get(target).doubleValue();
      if (dP == 0)
        dP = 0.00000000001;
      dH += Math.log(1.0 / dP) / Math.log(2.0);
      numActualTags++;
    }
    dH /= (double) numActualTags;
    dPP = Math.pow(2, dH);
    return dPP;
  }
}
